import torch
import torch.nn as nn
import cv2, os
from tqdm import tqdm
from glob import glob
import numpy as np
from torch.optim import Adam
from torchvision.utils import save_image
from random import randint
import torch.nn.functional as F

mse_loss = lambda x, y: F.mse_loss(x, y)


class Refiner(nn.Module):
    def __init__(self):
        super(Refiner, self).__init__()
        self.model = nn.Sequential(nn.Conv2d(3, 128, 3, stride=1, padding='same'), nn.ReLU(True),
                                   nn.Conv2d(128, 128, 3, stride=1, padding='same'), nn.ReLU(True),
                                   nn.Conv2d(128, 3, 3, stride=1, padding='same'), nn.Sigmoid())
    def forward(self, image):
        return self.model(image)
    

class Trainer():
    def __init__(self, root, scene, device):
        self.root = root
        self.scene = scene
        self.device = device
        self.model = Refiner().to(device)
        self.optimizer = Adam(self.model.parameters(), lr=0.00005)
        self.load_data()
        
    def load_data(self):
        train_path = os.path.join(self.root, self.scene, 'train', 'ours_30000', 'renders')
        gt_path = os.path.join(self.root, self.scene, 'train', 'ours_30000', 'gt')
        test_path = os.path.join(self.root, self.scene, 'video_images')
        
        train_img_names = sorted(glob(os.path.join(train_path, '*png')))
        gt_img_names = sorted(glob(os.path.join(gt_path, '*png')))
        test_img_names = sorted(glob(os.path.join(test_path, '*png')))
        
        train_images = []
        gt_images = []
        test_images = []
        for name in tqdm(train_img_names, desc=f'Loading {self.scene} train images ...'):
            train_images.append(cv2.imread(name) / 255.)
        for name in tqdm(gt_img_names, desc=f'Loading GT {self.scene} train images ...'):
            gt_images.append(cv2.imread(name) / 255.)
        for name in tqdm(test_img_names, desc=f'Loading test {self.scene} images ...'):
            test_images.append(cv2.imread(name) / 255.)
            
        self.train_images = torch.from_numpy(np.stack(train_images, 0)).permute(0, 3, 1, 2).float()
        self.gt_images = torch.from_numpy(np.stack(gt_images, 0)).permute(0, 3, 1, 2).float()
        self.test_images = torch.from_numpy(np.stack(test_images, 0)).permute(0, 3, 1, 2).float()
        
        assert self.train_images.shape[0] == self.gt_images.shape[0]
        print(f"Train data: {self.train_images.shape}")
        print(f"Test data: {self.test_images.shape}") 
        
    def batch_data(self, bs=32, ps=32):
        n_images = self.train_images.shape[0]
        idx = randint(0, n_images-1)
        train_img = self.train_images[idx]
        gt_img = self.gt_images[idx]
        H, W = train_img.shape[-2:]
        
        train_data, gt_data = [], []
        for i in range(bs):
            h = randint(0, H-ps)
            w = randint(0, W-ps)
            train_data.append(train_img[:, h: h+ps, w: w+ps])
            gt_data.append(gt_img[:, h: h+ps, w: w+ps])
        return torch.stack(train_data, 0).to(self.device), torch.stack(gt_data, 0).to(self.device)
        
    def train(self, iterations=20000):
        for it in tqdm(range(iterations), desc='Training ...'):
            n_images = self.train_images.shape[0]
            idx = randint(0, n_images-1)
            # train_data, gt_data = self.batch_data(bs=32, ps=32)
            train_data, gt_data = self.train_images[idx][None].to(self.device), self.gt_images[idx][None].to(self.device)
            pred = self.model(train_data)
            loss = mse_loss(pred, gt_data)
            
            if it % 1000 == 0:
                print(loss.item(), mse_loss(train_data, gt_data).item())
                
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
        self.test()
    
    def test(self):
        save_dir = os.path.join(self.root, 'refine_images')
        os.makedirs(save_dir, exist_ok=True)
        with torch.no_grad():
            for i, data in tqdm(enumerate(self.test_images), desc='Test ...'):
                data = data[None].to(self.device)
                pred = self.model(data).squeeze(0)
                save_image(pred, f'refined_{i:05d}.png')


if __name__ == '__main__':
    torch.cuda.set_device(5)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    root = 'output'
    for scene in ['bicycle', 'bonsai', 'counter', 'flowers', 'garden', 'kitchen', 'room', 'stump', "treehill"]:
        model = Trainer(root=root, scene=scene, device=device)
        model.train(iterations=20000)
    